#include "mpu6050.h"
#include <math.h>

/* ==== 内部寄存器 ==== */
#define RA_SMPLRT_DIV   0x19
#define RA_CONFIG       0x1A
#define RA_GYRO_CONFIG  0x1B
#define RA_ACCEL_CONFIG 0x1C
#define RA_ACCEL_XOUT_H 0x3B
#define RA_TEMP_OUT_H   0x41
#define RA_GYRO_XOUT_H  0x43
#define RA_PWR_MGMT_1   0x6B
#define RA_WHO_AM_I     0x75

/* 工具常量 */
#ifndef DEG_PER_RAD
#define DEG_PER_RAD 57.29577951308232f
#endif

/* ==== 内部 I2C 读写 ==== */
/* 说明：HAL 的 Mem_Read/Write 需要 8bit 地址，因此把 7bit 地址左移 1 */
static HAL_StatusTypeDef mpu_write_u8(MPU6050_t *dev, uint8_t reg, uint8_t val) {
  uint16_t addr8 = (dev->addr7 << 1);
  return HAL_I2C_Mem_Write(dev->hi2c, addr8, reg, I2C_MEMADD_SIZE_8BIT, &val, 1, 100);
}
static HAL_StatusTypeDef mpu_read(MPU6050_t *dev, uint8_t reg, uint8_t *buf, uint16_t len) {
  uint16_t addr8 = (dev->addr7 << 1);
  return HAL_I2C_Mem_Read(dev->hi2c, addr8, reg, I2C_MEMADD_SIZE_8BIT, buf, len, 100);
}

/* ==== 内部：设置灵敏度换算 ==== */
static void mpu_setup_sensitivity(MPU6050_t *dev) {
  switch (dev->acc_range) {
    case MPU6050_ACC_2G:  dev->acc_sens  = 16384.0f; break;
    case MPU6050_ACC_4G:  dev->acc_sens  = 8192.0f;  break;
    case MPU6050_ACC_8G:  dev->acc_sens  = 4096.0f;  break;
    default:              dev->acc_sens  = 2048.0f;  break; // 16G
  }
  switch (dev->gyro_range) {
    case MPU6050_GYRO_250DPS:  dev->gyro_sens = 131.0f;  break;
    case MPU6050_GYRO_500DPS:  dev->gyro_sens = 65.5f;   break;
    case MPU6050_GYRO_1000DPS: dev->gyro_sens = 32.8f;   break;
    default:                   dev->gyro_sens = 16.4f;   break; // 2000 dps
  }
}

/* ==== 内部：时间戳（us） ==== */
static inline uint32_t now_us(uint8_t use_dwt) {
#if (__CORTEX_M >= 0x03)
  if (use_dwt) {
    CoreDebug->DEMCR |= CoreDebug_DEMCR_TRCENA_Msk;
    DWT->CTRL |= DWT_CTRL_CYCCNTENA_Msk;
    return DWT->CYCCNT / (SystemCoreClock / 1000000U);
  }
#endif
  return HAL_GetTick() * 1000U; // 近似
}

/* ==== API 实现 ==== */
MPU6050_Result MPU6050_Init(MPU6050_t *dev,
                            I2C_HandleTypeDef *hi2c,
                            uint8_t addr7,
                            MPU6050_AccelRange acc_range,
                            MPU6050_GyroRange gyro_range,
                            MPU6050_DLPF dlpf,
                            uint8_t smplrt_div,
                            float alpha,
                            uint8_t use_dwt_time)
{
  dev->hi2c = hi2c;
  dev->addr7 = addr7;
  dev->acc_range = acc_range;
  dev->gyro_range = gyro_range;
  dev->alpha = (alpha <= 0.f || alpha >= 1.f) ? 0.98f : alpha;
  dev->pitch_deg = 0.f;
  dev->roll_deg  = 0.f;
  dev->use_dwt_time = use_dwt_time ? 1 : 0;

  dev->gyro_bias[0] = dev->gyro_bias[1] = dev->gyro_bias[2] = 0.f;
  dev->acc_bias[0]  = dev->acc_bias[1]  = dev->acc_bias[2]  = 0.f;

  mpu_setup_sensitivity(dev);

  HAL_Delay(100);

  /* 唤醒（PWR_MGMT_1=0x00 使用内部8MHz时钟或后续可设置PLL） */
  if (mpu_write_u8(dev, RA_PWR_MGMT_1, 0x00) != HAL_OK) return MPU6050_ERR;
  HAL_Delay(10);

  /* WHO_AM_I 检查（期望 0x68） */
  uint8_t who=0;
  if (mpu_read(dev, RA_WHO_AM_I, &who, 1) != HAL_OK) return MPU6050_ERR;
  if ((who & 0x7E) != 0x68) { // 某些芯片屏蔽位不同，这里放宽匹配
    return MPU6050_ERR;
  }

  /* 设置 DLPF + 采样率分频 */
  if (MPU6050_SetDlpfSmpl(dev, dlpf, smplrt_div) != MPU6050_OK) return MPU6050_ERR;

  /* 设置量程 */
  if (MPU6050_SetRanges(dev, acc_range, gyro_range) != MPU6050_OK) return MPU6050_ERR;

  /* 记录时间戳 */
  dev->last_us = now_us(dev->use_dwt_time);

  return MPU6050_OK;
}

MPU6050_Result MPU6050_SetRanges(MPU6050_t *dev,
                                 MPU6050_AccelRange acc_range,
                                 MPU6050_GyroRange gyro_range)
{
  /* Accel range */
  uint8_t acc_cfg = ((uint8_t)acc_range & 0x03) << 3; // AFS_SEL bits[4:3]
  if (mpu_write_u8(dev, RA_ACCEL_CONFIG, acc_cfg) != HAL_OK) return MPU6050_ERR;

  /* Gyro range */
  uint8_t gyro_cfg = ((uint8_t)gyro_range & 0x03) << 3; // FS_SEL bits[4:3]
  if (mpu_write_u8(dev, RA_GYRO_CONFIG, gyro_cfg) != HAL_OK) return MPU6050_ERR;

  dev->acc_range  = acc_range;
  dev->gyro_range = gyro_range;
  mpu_setup_sensitivity(dev);
  return MPU6050_OK;
}

MPU6050_Result MPU6050_SetDlpfSmpl(MPU6050_t *dev, MPU6050_DLPF dlpf, uint8_t smplrt_div)
{
  uint8_t cfg = ((uint8_t)dlpf & 0x07);
  if (mpu_write_u8(dev, RA_CONFIG, cfg) != HAL_OK) return MPU6050_ERR;
  if (mpu_write_u8(dev, RA_SMPLRT_DIV, smplrt_div) != HAL_OK) return MPU6050_ERR;
  return MPU6050_OK;
}

MPU6050_Result MPU6050_ReadRaw(MPU6050_t *dev, MPU6050_Raw *out)
{
  uint8_t buf[14];
  if (mpu_read(dev, RA_ACCEL_XOUT_H, buf, 14) != HAL_OK) return MPU6050_ERR;

  out->ax = (int16_t)((buf[0] << 8) | buf[1]);
  out->ay = (int16_t)((buf[2] << 8) | buf[3]);
  out->az = (int16_t)((buf[4] << 8) | buf[5]);
  out->temp = (int16_t)((buf[6] << 8) | buf[7]);
  out->gx = (int16_t)((buf[8] << 8) | buf[9]);
  out->gy = (int16_t)((buf[10] << 8) | buf[11]);
  out->gz = (int16_t)((buf[12] << 8) | buf[13]);
  return MPU6050_OK;
}

float MPU6050_ReadTempC(MPU6050_t *dev)
{
  uint8_t buf[2];
  if (mpu_read(dev, RA_TEMP_OUT_H, buf, 2) != HAL_OK) return NAN;
  int16_t raw = (int16_t)((buf[0] << 8) | buf[1]);
  /* 数据手册：Temp(°C) = raw/340 + 36.53 */
  return ((float)raw) / 340.0f + 36.53f;
}

MPU6050_Result MPU6050_CalibrateGyro(MPU6050_t *dev, uint16_t samples, uint16_t delay_ms)
{
  MPU6050_Raw r;
  double sx=0, sy=0, sz=0;

  for (uint16_t i=0; i<samples; ++i) {
    if (MPU6050_ReadRaw(dev, &r) != MPU6050_OK) return MPU6050_ERR;
    sx += r.gx;
    sy += r.gy;
    sz += r.gz;
    HAL_Delay(delay_ms);
  }
  dev->gyro_bias[0] = (float)(sx / samples) / dev->gyro_sens;
  dev->gyro_bias[1] = (float)(sy / samples) / dev->gyro_sens;
  dev->gyro_bias[2] = (float)(sz / samples) / dev->gyro_sens;
  return MPU6050_OK;
}

MPU6050_Result MPU6050_UpdateAngles(MPU6050_t *dev)
{
  MPU6050_Raw raw;
  if (MPU6050_ReadRaw(dev, &raw) != MPU6050_OK) return MPU6050_ERR;

  float acc_g[3], gyro_dps[3];
  MPU6050_ConvertAccel(dev, &raw, acc_g);   // ax, ay, az in g
  MPU6050_ConvertGyro (dev, &raw, gyro_dps); // gx, gy, gz in dps

  /* --- 关键修改：把“pitch”定义为绕 X 轴（roll-like）的角 --- */
  // 用 ay/az 计算：绕 X 轴的角（更稳建议用 atan2）
  float pitch_acc = atan2f(acc_g[1], acc_g[2]) * DEG_PER_RAD;  // ← 原来是用 ax/√(ay^2+az^2)

  /* 计算 dt */
  uint32_t t = now_us(dev->use_dwt_time);
  float dt = (t - dev->last_us) / 1e6f;
  if (dt <= 0.f || dt > 0.1f) dt = 0.01f;
  dev->last_us = t;

  /* --- 关键修改：绕 X 轴角速度用 gx --- */
  float pitch_gyro = dev->pitch_deg + gyro_dps[0] * dt;  // ← 保持使用 gx（X 轴）
  float roll_gyro  = dev->roll_deg  + gyro_dps[1] * dt;  // 这行随意保留/不用也行

  /* 互补滤波（alpha 已在 init 配好） */
  float a = dev->alpha;
  dev->pitch_deg = a * pitch_gyro + (1.f - a) * pitch_acc;

  // 你目前没用到 roll，可以保留原式：
  float roll_acc  = atanf(acc_g[1] / sqrtf(acc_g[0]*acc_g[0] + acc_g[2]*acc_g[2])) * DEG_PER_RAD;
  dev->roll_deg  = a * roll_gyro  + (1.f - a) * roll_acc;

  return MPU6050_OK;
}
